from django.utils import timezone
from datetime import timedelta
from tasks.models import Task
import pytz


def generate_time_series_data(mode, tz):
    now = timezone.now().astimezone(tz)

    if mode == 'daily':
        count = 7
        start_date = now - timedelta(days=count - 1)
        end_date = now
        period = 'day'

    elif mode == 'monthly':
        start_date = now.replace(month=1, day=1)
        end_date = now.replace(month=12, day=31)
        count = 12
        period = 'month'

    elif mode == 'yearly':
        current_year = now.year
        start_date = now.replace(year=current_year - 4, month=1, day=1)
        end_date = now.replace(month=12, day=31)
        count = 5
        period = 'year'

    else:
        raise ValueError("Invalid mode")

    labels = generate_date_labels(start_date, period, count)
    data = {
        'tasks_count': get_model_timeseries(Task, start_date, count, period, tz),
        'completed_tasks_count': get_model_timeseries(Task, start_date, count, period, tz),
        'todo_tasks_count': get_model_timeseries(Task, start_date, count, period, tz),
        'in_progress_tasks_count': get_model_timeseries(Task, start_date, count, period, tz),
        'overdue_tasks_count': get_model_timeseries(Task, start_date, count, period, tz),
    }
    return {'labels': labels, 'datasets': data}


def generate_date_labels(start_date, period, count):
    labels = []
    current = start_date

    if period == 'day':
        for _ in range(count):
            labels.append(current.strftime('%b %d').lstrip('0').replace(' 0', ' '))
            current += timedelta(days=1)

    elif period == 'month':
        for i in range(count):
            labels.append(f'{i + 1}月')

    elif period == 'year':
        for i in range(count):
            labels.append(str(start_date.year + i))

    return labels


def get_model_timeseries(model, start_date, count, period, tz):
    now = timezone.now().astimezone(tz)
    end_date = now

    utc_start = start_date.astimezone(pytz.utc)
    utc_end = end_date.astimezone(pytz.utc)

    results = (
        model.objects
        .filter(created_at__range=(utc_start, utc_end))
        .values_list('created_at', flat=True)
    )

    counts = [0] * count

    for dt in results:
        local_dt = dt.astimezone(tz)

        if period == 'day':
            diff = (local_dt.date() - start_date.date()).days
        elif period == 'month':
            diff = local_dt.month - 1
        elif period == 'year':
            diff = local_dt.year - start_date.year
        else:
            continue

        if 0 <= diff < count:
            counts[diff] += 1
    return counts
